#!/usr/bin/env python3
"""
基础聊天示例

这个示例展示了如何使用LangChain进行基本的对话。
演示了LLM的基本使用、消息处理和简单的对话流程。
"""

import sys
import os
from pathlib import Path

# 添加项目根目录到Python路径
project_root = Path(__file__).parent.parent
sys.path.append(str(project_root))

from langchain_core.messages import HumanMessage, AIMessage, SystemMessage
from rich.console import Console
from rich.panel import Panel
from rich.prompt import Prompt

from config.settings import settings
from core.llm_manager import llm_manager

console = Console()

def display_welcome():
    """显示欢迎信息"""
    welcome_text = """
🦜🔗 LangChain基础聊天示例

这个示例演示了：
1. 如何初始化LLM
2. 如何构建消息
3. 如何进行基本对话
4. 如何处理响应

输入 'quit' 退出程序
    """
    
    panel = Panel(
        welcome_text.strip(),
        title="基础聊天示例",
        border_style="blue",
        padding=(1, 2)
    )
    
    console.print(panel)
    console.print()

def basic_chat_example():
    """基础聊天示例"""
    console.print("[bold green]开始基础聊天示例...[/bold green]\n")
    
    # 检查LLM是否可用
    if not llm_manager.get_current_model():
        console.print("[red]错误: 没有可用的LLM模型[/red]")
        console.print("[yellow]请检查您的API密钥配置[/yellow]")
        return
    
    # 显示当前模型信息
    model_info = llm_manager.get_model_info()
    console.print(f"[cyan]当前模型: {model_info.get('name', '未知')}[/cyan]")
    console.print()
    
    # 系统提示词
    system_prompt = """你是一个友好的AI助手。请用中文回答问题，保持简洁和有用。"""
    
    # 对话历史
    conversation_history = [SystemMessage(content=system_prompt)]
    
    while True:
        try:
            # 获取用户输入
            user_input = Prompt.ask("[bold blue]您[/bold blue]").strip()
            
            if not user_input:
                continue
            
            if user_input.lower() in ['quit', 'exit', '退出']:
                console.print("[cyan]再见！[/cyan]")
                break
            
            # 添加用户消息
            user_message = HumanMessage(content=user_input)
            conversation_history.append(user_message)
            
            # 调用LLM
            console.print("\n[dim]正在生成响应...[/dim]")
            
            try:
                response = llm_manager.invoke(conversation_history)
                
                # 显示AI响应
                response_panel = Panel(
                    response,
                    title="AI助手",
                    border_style="green",
                    padding=(1, 2)
                )
                console.print(response_panel)
                
                # 添加AI响应到历史
                ai_message = AIMessage(content=response)
                conversation_history.append(ai_message)
                
            except Exception as e:
                console.print(f"[red]生成响应时出错: {str(e)}[/red]")
                console.print("[yellow]请检查您的API密钥和网络连接[/yellow]")
            
            console.print()
        
        except KeyboardInterrupt:
            console.print("\n[cyan]程序已退出[/cyan]")
            break
        
        except Exception as e:
            console.print(f"[red]发生错误: {str(e)}[/red]")

def streaming_chat_example():
    """流式聊天示例"""
    console.print("[bold green]开始流式聊天示例...[/bold green]\n")
    
    if not settings.ENABLE_STREAMING:
        console.print("[yellow]流式输出未启用，跳过此示例[/yellow]")
        return
    
    # 简单的流式对话
    messages = [
        SystemMessage(content="你是一个有用的AI助手。"),
        HumanMessage(content="请用中文介绍一下LangChain框架，大约100字。")
    ]
    
    console.print("[cyan]流式响应演示:[/cyan]")
    console.print("[dim]AI助手: [/dim]")
    
    try:
        response_text = ""
        for chunk in llm_manager.stream(messages):
            # 使用 sys.stdout 来实现流式输出
            sys.stdout.write(chunk)
            sys.stdout.flush()
            response_text += chunk
        
        console.print("\n")
        
        # 显示完整响应
        if response_text:
            response_panel = Panel(
                response_text,
                title="完整响应",
                border_style="green",
                padding=(1, 2)
            )
            console.print(response_panel)
    
    except Exception as e:
        console.print(f"\n[red]流式响应出错: {str(e)}[/red]")

def message_types_example():
    """消息类型示例"""
    console.print("[bold green]消息类型示例...[/bold green]\n")
    
    # 演示不同类型的消息
    messages = [
        SystemMessage(content="你是一个专业的编程助手。"),
        HumanMessage(content="什么是Python？"),
        AIMessage(content="Python是一种高级编程语言，以其简洁和可读性著称。"),
        HumanMessage(content="请给我一个简单的Python代码示例。")
    ]
    
    console.print("[cyan]消息历史:[/cyan]")
    for i, message in enumerate(messages, 1):
        message_type = message.__class__.__name__
        content = message.content[:50] + "..." if len(message.content) > 50 else message.content
        console.print(f"{i}. [{message_type}] {content}")
    
    console.print("\n[cyan]AI响应:[/cyan]")
    
    try:
        response = llm_manager.invoke(messages)
        response_panel = Panel(
            response,
            title="AI助手",
            border_style="green",
            padding=(1, 2)
        )
        console.print(response_panel)
    
    except Exception as e:
        console.print(f"[red]生成响应时出错: {str(e)}[/red]")

def main():
    """主函数"""
    console.print("[bold cyan]🦜🔗 LangChain基础聊天示例[/bold cyan]\n")
    
    # 验证配置
    if not settings.validate_configuration():
        console.print("[red]配置验证失败，请检查您的设置[/red]")
        return
    
    display_welcome()
    
    # 运行示例
    examples = [
        ("消息类型示例", message_types_example),
        ("流式聊天示例", streaming_chat_example),
        ("交互式聊天", basic_chat_example)
    ]
    
    for name, example_func in examples:
        console.print(f"\n{'='*50}")
        console.print(f"运行: {name}")
        console.print('='*50)
        
        try:
            example_func()
        except Exception as e:
            console.print(f"[red]示例运行失败: {str(e)}[/red]")
        
        # 询问是否继续
        if name != examples[-1][0]:  # 不是最后一个示例
            if not Prompt.ask("\n继续下一个示例？", choices=["y", "n"], default="y") == "y":
                break
    
    console.print("\n[cyan]示例演示完成！[/cyan]")

if __name__ == "__main__":
    main() 